from inference import Inference
import json
import requests
import arguments as args

predictor1 = Inference(args.OPERATE_MODEL_CONFIG)
print('模型加载成功.....')


def find_substring_index(text, sub_text):
    return text.find(sub_text)


operate_type_lists = [{"operate_type": "1", "schema": ["任务编号"]}, {"operate_type": "2", "schema": ["工作站名称"]},
                      {"operate_type": "3", "schema": ["订单编号"]}, {"operate_type": "4", "schema": ["产品"]},
                      {"operate_type": "5", "schema": ["优先级"]}, {"operate_type": "6", "schema": ["SO交付时间"]},
                      {"operate_type": "9", "schema": ["订单编号", "TSK排程结束时间"]},
                      {"operate_type": "11", "schema": ["TASK状态"]}, {"operate_type": "12", "schema": ["MO状态"]},
                      {"operate_type": "13", "schema": ["TSK排程开始时间"]},
                      {"operate_type": "14", "schema": ["TSK排程结束时间"]},
                      {"operate_type": "15", "schema": ["TSK实际开始时间"]},
                      {"operate_type": "16", "schema": ["TSK实际结束时间"]},
                      {"operate_type": "17", "schema": ["TSK延期时间"]},
                      {"operate_type": "19", "schema": ["MO重排前预计交付时间"]},
                      {"operate_type": "20", "schema": ["MO重排后预计交付时间"]},
                      {"operate_type": "21", "schema": ["MO重排差异"]},
                      {"operate_type": "22", "schema": ["MO延期时间"]}]


def get_operate_type(xxx):
    for operate_type in operate_type_lists:
        if operate_type['operate_type'] == xxx:
            return operate_type['schema'][0]

    return '未定义'


if __name__ == '__main__':

    with open('./data/file.txt', 'r') as file:
        lines = file.readlines()
        print(f'ssssssssssss{lines}')
        list = []
        list_no = []
        for text in lines:
            print(text)
            text = json.loads(text)
            text = text['text']
            # text = "能给我展示一下TASK202400000501的生产任务详情。"
            text = text.replace('\n', '')
            operate_type_list = predictor1.predict(text)
            operate_type = operate_type_list[0]
            cats_tag = operate_type['label']
            score = operate_type['score']
            api = 'http://127.0.0.1:8892/ner'
            params = {
                'query': text,
                'operate_type': cats_tag
            }
            cats = get_operate_type(cats_tag)
            if (str(cats_tag) != '0') & (score > 0.15):
                response = requests.post(api, json=params).json()
                item = {"text": text, "cats": cats}
                item.update(response)
                list.append(item)
            else:
                print(f'无操作类型text:{text},operate_type_list:{operate_type_list}')
                list_no.append({"text": text, "cats": '未定义'})
        with open('./data/automatic_standards.jsonl', 'w', encoding='utf-8') as json_file:
            for item in list:
                json_file.write(json.dumps(item, ensure_ascii=False) + '\n')
            for item in list_no:
                json_file.write(json.dumps(item, ensure_ascii=False) + '\n')
        print("操作完成")
